#!/usr/local/bin/python
#-*- coding: utf-8 -*-
#############################################################################
#                                                                           #
#   File: fetch.py                                                          #
#                                                                           #
#   Copyright (C) 2008-2011 lyricconch <lyricconch@gmail.com>               #
#                                                                           #
#   This file is part of WallProxyMod.                                      #
#                                                                           #
#   WallProxyMod is free software: you can redistribute it and/or modify    #
#   it under the terms of the GNU General Public License as                 #
#   published by the Free Software Foundation, either version 3 of the      #
#   License, or (at your option) any later version.                         #
#                                                                           #
#   WallProxyMod is distributed in the hope that it will be useful,         #
#   but WITHOUT ANY WARRANTY; without even the implied warranty of          #
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the           #
#   GNU General Public License for more details.                            #
#                                                                           #
#   You should have received a copy of the GNU General Public License       #
#   along with WallProxyMod.  If not, see <http://www.gnu.org/licenses/>.   #
#                                                                           #
#############################################################################

from google.appengine.runtime import apiproxy_errors
from google.appengine.api import apiproxy_stub_map
from google.appengine.api import urlfetch_service_pb
from google.appengine.ext import db
from google.appengine.api import memcache
from google.appengine.api import users

from cStringIO import StringIO
from Crypto.Cipher import ARC4
import rsaasn1 

import sys, os, binascii, hashlib, pickle
import logging, traceback, re, pprint
import gzip, time, zlib

__author__ = "lyricconch@gmail.com"
__version__ = "2.0.1a"

def POST(rfile, wfile, env):
    time0 = time.time()
    SESSION = binascii.a2b_base64(env["HTTP_X_AUTHORIZATION"])
    SESSION = CONFIG.SITE_CIPHER.decrypt(SESSION)
    rpcs, request = [], rfile.read()
    if hashlib.sha512(request).digest() != SESSION[64:128]:
        logging.warn("Bad Request: digest %s, data: %s...",
                      binascii.b2a_hex(SESSION[64:128]), request[:100])
        print "Status: 400 Request is broken"
        return
    request = ARC4.new(SESSION[0:64]).decrypt(request)
    if CONFIG.TRANS_PICKLE:
        for urlfetch in pickle.loads(request):
            rpcs.append(make_fetch(*urlfetch, **CONFIG.FETCH_FLAGS))
        response = [rpc.get_result() for rpc in rpcs]
        response = pickle.dumps(response, pickle.HIGHEST_PROTOCOL)
    else:
        input, output = StringIO(request), StringIO()
        line = input.readline()
        while line.strip():
            method, path, ver = line.split()
            line, headers = input.readline(), dict()
            while line.strip():
                name, value = line.split(":", 1)
                headers[name.title()] = value.strip()
                line = input.readline()
            content = input.read(int(headers.get("Content-Length", 0)))
            line = input.readline()
            rpc = make_fetch(method, path, headers, content, **CONFIG.FETCH_FLAGS)
            rpcs.append(rpc)
        for rpc in rpcs:
            code, message, headers, content = rpc.get_result()
            output.write("%s %s %s\r\n" % (ver, code, message))            for header in headers.items():
                output.write("%s: %s\r\n" % header)
            output.write("\r\n")
            output.write(content)
        response = output.getvalue()
    response = ARC4.new(SESSION[0:64]).encrypt(response)
    SESSION = hashlib.sha512(request).digest() + hashlib.sha512(response).digest()
    SESSION = binascii.b2a_base64(SESSION).strip()
    print "Content-Type: %s" % CONFIG.TRANS_CTYPE
    print "X-Validation: %s" % SESSION
    print
    wfile.write(response)
    logging.info("finish request from %s, time %s, in %s, out %s", 
                 env["REMOTE_ADDR"], time.time()- time0, len(request), len(response))

def GET(rfile, wfile, env):
    print "Content-Type: text/plain"
    print
    print pprint.pformat(env)

def make_fetch(method, path, headers={}, payload="",
               message=None, deadline=None, redirect=None, validate=None):
    rpc = apiproxy_stub_map.UserRPC('urlfetch', CONFIG.FETCH_DEADLINE)
    request = urlfetch_service_pb.URLFetchRequest()
    response = urlfetch_service_pb.URLFetchResponse()
    request.set_url(path.encode("UTF-8") if isinstance(path, unicode) else path)
    request.set_method(getattr(request, method))
    if payload and method in ("POST", "PUT"):
        request.set_payload(payload)
    for key, value in headers.iteritems():
        if ishopbyhop(key): continue
        header = request.add_header()
        header.set_key(key)
        header.set_value(str(value))
    if deadline is None: deadline = CONFIG.FETCH_DEADLINE * (0.4 if message else 0.9)
    request.set_deadline(deadline)
    if redirect is None: redirect = method == "GET" and "cookie" not in headers
    request.set_followredirects(redirect)
    if validate is None: validate = method in "GET,HEAD" and "cookie" not in headers
    request.set_mustvalidateservercertificate(validate)
    rpc.make_call('Fetch', request, response, hook_result,
                  [method, path, headers, payload, message])
    return rpc

_re_setcookie = re.compile(r", (?=[^,;]+=)")
_re_contentrange = re.compile(r"byte (\d+)-(\d+)/(\d+)")
def hook_result(rpc):
    args = rpc.user_data
    try: rpc.check_success()
    except apiproxy_errors.ApplicationError, e:
        etype = urlfetch_service_pb.URLFetchServiceError
        errno = e.application_error
        code, reason = 440, "E\tUrlfetch: " + etype.ErrorCode_Name(errno)
        logging.warn("%s at %s %s", reason, args[0], args[1])
        if errno == etype.RESPONSE_TOO_LARGE and args[0] == "GET" and args[4][0] != "C":
            args[2]["range"] = "bytes=0-100000"
            args[4] = "C\t" + reason
            try: return make_fetch(*rpc.user_data, **CONFIG.FETCH_FLAGS).get_result()
            except: logging.warn("RangeRetry: %s at %s %s \n %s",
                                 sys.exc_info()[1], args[0], args[1], pprint.pformat(args[2]))
    except apiproxy_errors.Error, e:
        code, reason = 540, "E\tApiproxy: " + e.__class__.__name__
        logging.error("%s at %s %s", reason, args[0], args[1])
    except:
        code, reason = 541, "E\tUnhanlded: %s" % sys.exc_info()[1]
        logging.critical("%s at %s %s", reason, args[0], args[1], exc_info=True)
    else:
        ret = rpc.response
        code, content = str(ret.statuscode()), ret.content()
        message, headers = args[4] or "S\t", dict()
        for h in ret.header_list():
            key, value = h.key(), h.value()
            if ishopbyhop(key): continue
            headers[key.title()] = str(value)            
        if ret.finalurl(): 
            message += "#R" + ret.finalurl()
        if "Set-Cookie" in headers:
            headers["Set-Cookie"], count = _re_setcookie.subn("\r\nSet-Cookie: ", headers["Set-Cookie"])
            message += "#C%d" % (count + 1)
        if ret.contentwastruncated():
            message += "#T"
            if args[0] == "GET" and code in "200,203,206":
                code, message = "206", "C" + message[1:]
                s, _, t = (_re_contentrange.match(headers["Content-Range"]).groups()
                    if "Content-Range" in headers else (0, 0, headers.get("Content-Length", "*")))
                headers["Content-Range"] = "byte %s-%s/%s" % (s, s + len(content) - 1, t)
        headers['Content-Length'] = len(content)
        logging.debug("done: %s %s for %s %s", code, message, args[0], args[1])
        return code, message, headers, content
    return code, reason, {}, ""

_tp_hoppish = ("connection", "host", "keep-alive", "proxy-connection",
             "te", "trailers", "transfer-encoding", "upgrade",)
def ishopbyhop(name):
    return name.lower() in _tp_hoppish

def main():
    global CONFIG
    try:
        if not CONFIG:
            ID = os.environ["APPLICATION_ID"] + os.environ["PATH_INFO"]
            CONFIG = memcache.get(key=ID, namespace="FETCH_CONFIG")
            if not CONFIG:
                CONFIG = FETCH_CONFIG.all().filter("SITE_ID =", ID).get()
                if not CONFIG:
                    CONFIG = FETCH_CONFIG(SITE_ID=ID)
                    CONFIG.put()
                CONFIG.SITE_CIPHER = rsaasn1.importKey(open(CONFIG.SITE_KEY).read())
                _ = CONFIG.FETCH_PARAMS
                CONFIG.FETCH_FLAGS = dict(deadline=float(_.split(",")[0]) if _[0].isdigit() else None,
                                          redirect=0 if "~F" in _ else 1 if "F" in _ else None,
                                          validate=0 if "~V" in _ else 1 if "V" in _ else None)
            memcache.set(key=ID, value=CONFIG, namespace="FETCH_CONFIG")
        method = os.environ['REQUEST_METHOD']
        dispatcher = method.isupper() and globals().get(method)
        if not dispatcher:
            print "Status: 405 Method not Allow"
            return
        dispatcher(sys.stdin, sys.stdout, os.environ)
    except:
        logging.critical("Exception catch by main(), Environ: \n%s",
                         pprint.pformat(os.environ), exc_info=True)
        sys.stdout.reset()
        print "Status: 500 Internal Server Error"
        print "Content-Type: text/plain"
        print
        if CONFIG is None or CONFIG.SITE_DEBUG:
            traceback.print_exc(None, sys.stdout)

class FETCH_CONFIG(db.Model):
    SITE_ID = db.ByteStringProperty(required=True)
    SITE_KEY = db.ByteStringProperty(default="fetch.key")
    SITE_DEBUG = db.BooleanProperty(default=True)
    FETCH_DEADLINE = db.IntegerProperty(default=10)
    FETCH_PARAMS = db.ByteStringProperty(default=",,V")
    FETCH_RANGE = db.ByteStringProperty(default="bytes=0-100000")
    TRANS_PICKLE = db.BooleanProperty(default=False)
    TRANS_CTYPE = db.ByteStringProperty(default="application/octet-stream")
    TRANS_ZIPPED = db.ByteStringProperty(default="text/html;")
    CACHE_MODE = db.ByteStringProperty(default="NONE")
    CACHE_HEADER = db.ByteStringProperty(default="NONE")
    CACHE_SIZEMAX = db.IntegerProperty(default=0)
    CACHE_EXPIRE = db.IntegerProperty(default=0)

CONFIG = None
if __name__ == "__main__":
    main()
